import xml.etree.ElementTree as ET;
from Activity import Activity;
import sys, math;

class InputSim(object):
    def __init__(self, pgm, sStatesFile, sTransFile, sPriorsFile, bProbs_ByAct=False):
        self.pgm = pgm;
        self.nObsId = 0;
        self.bProbs_ByAct = bProbs_ByAct;
        [nStatesT0, nStatesT1] = self.getStateSpace_Sizes(sStatesFile);  #@UnusedVariable
        self.readTransitionM(sTransFile);
        self.readPriors(sPriorsFile);

    def getStateSpace_Sizes(self, sStatesFile):
        self.sActivities = dict();
        #Template of All Possible Arg Combinations (for ANY activity having 1, 2 or 3 args) in T0 & T1 respectively
        sArgsT0 = [[["A1"]], [["A1", "*"], ["*", "A1"], ["A1", "A2"], ["A2", "A1"]],
                   [["A1", "*", "*"], ["*", "A1", "A2"], ["A1", "A2", "A3"], ["A2", "A3", "A1"], ["A2", "A1", "A3"]],
                   [["-", "-", "-", "-"]]];
        sArgsT1 = [[["A1"]], [["A1", "*"], ["*", "A1"], ["A1", "A2"], ["A2", "A1"], ["A1", "X1"], ["X1", "A1"]], 
                   [["A1", "*", "*"], ["A1", "A2", "A3"], ["A2", "A3", "A1"], ["A2", "A1", "A3"], ["A1", "X1", "X2"], ["X1", "A1", "X2"], ["X1", "X2", "A1"]],
                   [["-", "-", "-", "-"]]];
        self.aAllArgs = [sArgsT0, sArgsT1];
        #self.aActNames = ["Null", "Loitering", "Wait-N-Meet", "Chase"]; #, "Approach"];
#         self.aActNames = ["Null", "Run", "Chase"]; #, "Approach"];
        self.aActNames, nArgs = self.readStateSpace(sStatesFile);
#         print self.aActNames, nArgs;
#         nArgs = [1,1,2,2,2];
#         nArgs = [1,1,2];
        aActSizeT0, aActSizeT1 = [1, 4, 5, 1], [1, 6, 7, 1]; #Total Arg Combinations (for ANY activity having 1, 2 or 3 args) in T0 & T1 respectively        
        iT0, iT1 = 0, 0;
        for i, sAct in enumerate(self.aActNames):
            if not(self.bProbs_ByAct): #If we consider the order of Args from Activity's Transitions & Priors, then use Template's Indices
                self.sActivities[sAct] = [nArgs[i], iT0, iT1]; #nArgs, Act_Index(t), Act_Index(t+1)
                iT0 += aActSizeT0[nArgs[i]-1]; #Accumulating all total arg combinations on T0 (from ActNames given)
                iT1 += aActSizeT1[nArgs[i]-1]; #Accumulating all total arg combinations on T1 (from ActNames given)
            else: #Else If we ignore the order of Args from Activity's Transitions & Priors, then Ignore Templates Indices
                self.sActivities[sAct] = [nArgs[i], iT0, iT1]; #nArgs, Act_Index(t), Act_Index(t+1)
                iT0 += 1; #Set Number of Arg combinations on T0 (from ActNames given) as only One
                iT1 += 1; #Set Number of Arg combinations on T1 (from ActNames given) as only One
        return [iT0, iT1];    
            
    def getIndex(self, aActT0, aActT1):
        mArg = dict(); #Mapping from Arg_Input -> Arg_Template 
        unkArgs = []; #Args from ActT0 that Didn't Match
        #Search for Args Matching on both Activities
        k = 1;
        for argT0 in aActT0[1]:
            if (argT0!="*") and (argT0 in aActT1[1]):
                mArg[argT0] = "A"+str(k); #Args from ActT0 that Did Match
                k += 1;
            elif (argT0!="*"):
                unkArgs.append(argT0);
        if len(mArg)==0:
            return [None, None];
        #Map remaining Args (that Didn't Match) with new Arg_Template
        for i, unkArg in enumerate(unkArgs, len(mArg)+1):
            mArg[unkArg] = "A"+str(i); #Args from ActT0 that Didn't Match
        k = 1;
        for argT1 in aActT1[1]:
            if (argT1!="*") and not(argT1 in aActT0[1]):
                mArg[argT1] = "X"+str(k); #Args from ActT1 that Didn't Match
                k += 1;
        #Build Acts with Arg_Templates
        aArgsT0, aArgsT1 = [], [];
        for argT0 in aActT0[1]:
            if argT0=="*": aArgsT0.append("*");
            else: aArgsT0.append(mArg[argT0]);
        for argT1 in aActT1[1]:
            if argT1=="*": aArgsT1.append("*");
            else: aArgsT1.append(mArg[argT1]);
        #Search for their Index (for Transition Matrix)
        sActT0, sActT1 = aActT0[0][0], aActT1[0][0];
#        print sActT0, aArgsT0;
#        print sActT1, aArgsT1;
        nArgsT0 = self.sActivities[sActT0][0]-1; #nArgs of this Activity T0
        iT0, iT1 = -1, -1;
        if self.bProbs_ByAct and len(aArgsT0)>0 and len(aArgsT1)>0:
            iT0 = self.sActivities[sActT0][1];
            iT1 = self.sActivities[sActT1][2];
            return [iT0, iT1];
        if aArgsT0 in self.aAllArgs[0][nArgsT0]:
            iT0 = self.sActivities[sActT0][1];
            iT0 += self.aAllArgs[0][nArgsT0].index(aArgsT0);
        nArgsT1 = self.sActivities[sActT1][0]-1; #nArgs of this Activity T0
        if aArgsT1 in self.aAllArgs[1][nArgsT1]: 
            iT1 = self.sActivities[sActT1][2];
            iT1 += self.aAllArgs[1][nArgsT1].index(aArgsT1);
#        print self.sActivities[self.aActNames[-1]][1]
        if iT0<0 and iT1<0:
            print "---\nError: Match (on Transition table) not found for activities:", aActT0, "and", aActT1;
        elif iT0<0:
            print "---\nError: Match (on Transition table) not found for activity T0: ", aActT0, aActT1, mArg;
        elif iT1<0:
            print "---\nError: Match (on Transition table) not found for activity T1: ", aActT1, aActT0, mArg;
#         if iT0<0 or iT1<0:
#             sys.exit();
        return [iT0, iT1];
    
    def getIndexPrior(self, aActT0):
        mArg = dict(); #Mapping from Arg_Input -> Arg_Template
        #Search for Args in this Activity
        k = 1;
        for argT0 in aActT0[1]:
            if (argT0!="*"):
                mArg[argT0] = "A"+str(k); #Args from ActT0 that Did Match
                k += 1;
        #Build Acts with Arg_Templates
        aArgsT0 = [];
        for argT0 in aActT0[1]:
            if argT0=="*": aArgsT0.append("*");
            else: aArgsT0.append(mArg[argT0]);
        #Search for their Index (for Prior Matrix)
#         print "*****", aActT0, str(aActT0[0][0]);
        sActT0 = aActT0[0][0];
#        print sActT0, aArgsT0;
        nArgsT0 = self.sActivities[sActT0][0]-1; #nArgs of this Activity T0
        iT0 = -1;
        if self.bProbs_ByAct and len(aArgsT0)>0:
            iT0 = self.sActivities[sActT0][1];
            return iT0;
        if aArgsT0 in self.aAllArgs[0][nArgsT0]:
            iT0 = self.sActivities[sActT0][1];
            iT0 += self.aAllArgs[0][nArgsT0].index(aArgsT0);
        if iT0<0:
            print "---\nError: Match (on Priors table) not found for activity: ", aActT0;
            sys.exit();
        return iT0;
        
    def makeXmlActivities(self, sXmlFile=""):
        xmlAct = ET.parse(sXmlFile);
        actCollection = xmlAct.getroot();
        actElements = actCollection.findall("activity");
        self.aActivities = []; #All Activities read from XML
        for actItem in actElements:
            #Activity:
            if self.createAct(actItem, bRoot=True):
                nObs = len(actItem.findall("observable"));
#                 self.createObs(self.aActivities[-1], nObs);
                #Check Observations:
#                 obsActElements = actItem.findall("observable");                
#                 for obsActItem in obsActElements:
#                     #Observation:
#                     self.createAct(obsActItem, bRoot=False, activity=self.aActivities[-1]);
# #         print "---Activity Nodes from XML:";
# #         for actNode in self.aActivities:
# #             self.printActNode(actNode);
                
    def createAct(self, actItem, bRoot=False, activity=None):
        aLatestActors = []; #Only Latest Actors observed in this Event
        aActorsAll = []; #All Actors seen so far in this Event
        aActorTypes = []; #All Actor Types
        aActorTimes = []; #All Actor Times
        aActorRoles = []; #All Actor Roles
        aActorsInfo = []; #All Actor Info (ActorsAll, Types, Times, Roles)
        actorElements = actItem.findall("actor");
        for actorItem in actorElements:
            bActorObserved = (actorItem.get("observed")=="True");
            sName = actorItem.get("name") if bActorObserved else "*";
            aActorsAll.append(sName);
            aActorTypes.append(actorItem.get("type"));
            aActorRoles.append(actorItem.get("role"));
            aActorTimes.append(str(actorItem.get("time").replace("[","").replace("]","")).split(","));            
        aEndT = [];
        for aTime in aActorTimes:
            aEndT.append(int(aTime[-1].replace("+","")));
        nEndT = max(aEndT);
        for i,actor in enumerate(aActorsAll):
            if nEndT==int(aActorTimes[i][-1].replace("+","")) and not(actor in aLatestActors):
                aLatestActors.append(actor);
        aActorsInfo.append(aActorsAll);
        aActorsInfo.append(aActorTypes);
        aActorsInfo.append(aActorTimes);
        aActorsInfo.append(aActorRoles);
        sActName = actItem.get("name");
         
        aTimes = str(actItem.get("time").replace("[","").replace("]","")).split(",");
        nRealInitT = int(aTimes[0]);
        nTime = int(aTimes[-1].replace("t","").replace("+",""));
        if bRoot==True:
            sId = actItem.get("id") + "." + str(nTime);
        else:
            sId = activity.rootAct.sId.split(".")[0] + "_Obs" + str(self.nObsId) + "." + str(activity.rootAct.aTimes[-1].replace("+",""));
            self.nObsId += 1;
        sParentsId = [];
        if bRoot==False:
            sParentsId = str(actItem.get("parents").replace("[","").replace("]","")).split(",");
            for i in range(len(sParentsId)):
                sParentsId[i] += "." + str(nTime);
        if len(sParentsId)==1 and sParentsId[0]=="": sParentsId = [];
        sScores = [];
        if actItem.get("score")!=None:
            sScores = str(actItem.get("score").replace("[","").replace("]","")).split(",");
        aScores = [];
        for strScore in sScores:
            if strScore!="" and strScore!="nan":
                if float(strScore)<0: #If it's a Log-Likelihood
                    aScores.append(math.exp(float(strScore))); #REMOVE LOG-LIKELIHOODS!!!
        if len(aScores)==0:
            scoreElements = actItem.findall("score");
            for scoreItem in scoreElements:
                strScore = scoreItem.get("value");
                if strScore!="" and strScore!="nan":
                    if float(strScore)<=0: #If it's a Log-Likelihood
                        aScores.append(math.exp(float(strScore))); #REMOVE LOG-LIKELIHOODS!!!
                    elif float(strScore)>0: #If it's a Likelihood
                        aScores.append(float(strScore));
        sLocation = actItem.get("location");
        #Check whether to ignore the whole Activity/Sub-Activity
        try:
            bNullPersons = False;
            for sName in aActorsAll:
                if sName!="*":
                    bNullPersons = False;
            if nTime<0 or aScores==[] or bNullPersons:
                raise SystemExit;
        except SystemExit:
            if nTime<0:
                print "---\nError: Negative Times were found in Activity:", sActName, aTimes;
            elif aScores==[]:
                print "---\nError: Score not found in Activity:", sActName;                
            elif bNullPersons:
                print "---\nError: Persons not found in Activity:", sActName, aActorsAll;
            sys.exit();
        #Add the Activity or Sub-Activity
        if bRoot: 
            activity = Activity(self.pgm, sActName, aLatestActors, nTime, sId);
            self.aActivities.append(activity);
        activity.makeXmlNode(sId, sActName, aLatestActors, aActorsInfo, aTimes, 
                             nRealInitT, sParentsId, aScores, sLocation, bRoot);
        return True;
    
    def createObs(self, activity, nObs):
        sId = activity.rootAct.sId.split(".")[0] + "_Obs" + "." + str(activity.rootAct.aTimes[-1]);
        sActName = activity.rootAct.actName;
        aLatestActors = activity.rootAct.aSubjects;
        aActorsInfo = [];
        aTimes = activity.rootAct.aTimes;
        nRealInitT = 0;
        sParentsId = [activity.rootAct.sId];
        aScores = activity.rootAct.aScores;
        sLocation = activity.rootAct.location;
#         activity.nProb_ObsNull = activity.nProb_ObsNull**nObs;
#         if nObs==1: activity.nProb_ObsNull=0.1;
#         elif nObs==4: activity.nProb_ObsNull=0.04;
        activity.makeXmlNode(sId, sActName, aLatestActors, aActorsInfo, aTimes, 
                             nRealInitT, sParentsId, aScores, sLocation, bRoot=False);
        
    def printActNode(self, actNode):
        print "Activity Node:";
        print "sName:", actNode.rootAct.name;
        print "sId:", actNode.rootAct.sId;
        print "sActName:", actNode.rootAct.actName;
        print "aSubjects:", actNode.rootAct.aSubjects;
        print "aActorsAll:", actNode.rootAct.aActorsAll;
        print "aActorTypes:", actNode.rootAct.aActorTypes;
        print "aActorTimes:", actNode.rootAct.aActorTimes;
        print "aActorRoles:", actNode.rootAct.aActorRoles;
        print "aTimes:", actNode.rootAct.aTimes;
        print "aScores:", actNode.rootAct.aScores;
        print "sLocation:", actNode.rootAct.location;
        print "sParentsId:", actNode.rootAct.sParentsId;
        print "";
        for subAct in actNode.aChildNodes:
            print "- Observable Node:";
            print "- sName:", subAct.name;
            print "- sId:", subAct.sId;
            print "- sObsName:", subAct.actName;
            print "- aSubjects:", subAct.aSubjects;
            print "- aActorsAll:", subAct.aActorsAll;
            print "- aActorTypes:", subAct.aActorTypes;
            print "- aActorTimes:", subAct.aActorTimes;
            print "- aActorRoles:", subAct.aActorRoles;
            print "- aTimes:", subAct.aTimes;
            print "- aScores:", subAct.aScores;
            print "- sLocation:", subAct.location;
            print "- sParentsId:", subAct.sParentsId;
            print "";

    def readTransitionM(self, sTransFile):
        self.aTrans = [];
        f = open(sTransFile, "r");
        while 1:
            sLine = f.readline();
            aLine = sLine.split(",");
            for i,sProb in enumerate(aLine):
                sProb = sProb.replace("\n","");
                if sProb.replace(".","").isdigit():
                    aLine[i] = float(sProb);
            if not sLine:
                break;
            self.aTrans.append(aLine);
        f.close();
#         for i in range(len(self.aTrans)):
#             print self.aTrans[i];

    def readPriors(self, sPriorsFile):
        f = open(sPriorsFile, "r");
        while 1:
            sLine = f.readline();
            aLine = sLine.split(",");
            for i,sProb in enumerate(aLine):
                sProb = sProb.replace("\n","");
                if sProb.replace(".","").isdigit():
                    aLine[i] = float(sProb);
            if not sLine:
                break;
            self.aPriors = aLine;
        f.close();
#         print self.aPriors;

    def readStateSpace(self, sStatesFile):
        aActNames = [];
        nArgs = [];
        f = open(sStatesFile, "r");
        while 1:
            sLine = f.readline();
            sLine = sLine.replace("\n","");
            if not sLine:
                break;
            aLine = sLine.split(",");
            aActNames.append(aLine[0]);
            nArgs.append(int(aLine[1]));
        return aActNames, nArgs;
            
            
    def printStateSpace(self):
        for i in range(len(self.aStates)):
            print i+1, "-", self.aStates[i];
        